package com.RJAI.gemini.service.impl;

import cn.hutool.core.util.RandomUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.RJAI.gemini.utils.GCloudUtils;
import com.RJAI.gemini.exception.RJException;

import okhttp3.OkHttpClient;
import okhttp3.Request;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.TimeUnit;
/*
作者信息
WeChat:ciansec
QQ:1068092013
 */
@Service
public class ClaudeServiceImpl {
    @Resource
    private RedisTemplate redisTemplate;
    @Autowired
    private HttpServletRequest httpServletRequest;


    @Value("${cian.serviceAccountDir}")
    private String serviceAccountDir;

    @Value("${cian.serviceAccountExpiration}")
    private Integer serviceAccountExpiration;

    @Autowired
    private AsyncService asyncService;

    /**
     * 处理流式对话
     * @param response
     * @param key
     * @param jsonObject
     * @param model
     * @return
     * @throws IOException
     */
    public Object claudeMessageWithStream(HttpServletResponse response,String key, JSONObject jsonObject, String model) throws IOException {
        jsonObject.put( "anthropic_version", "vertex-2023-10-16" );
        //用来拼接请求头
        String region = null;
        if (model.contains( "claude-3-5-sonnet" )){
            switch (RandomUtil.randomInt( 1,3 )){
                case 1:
                    region="europe-west1";
                    break;
                case 2:
                    region="us-east5";
                    break;
                default:
                    region="europe-west1";

            }
        }else if (model.contains("claude-3-opus")){
            switch (RandomUtil.randomInt( 1,3 )){
                case 1:
                    region="europe-west1";
                    break;
                case 2:
                    region="us-east5";
                    break;
                default:
                    region="europe-west1";
            }
        }





        //下面代码用来调用服务获取accessToken和project_id
        String accessToken = null;
        String project_id = null;
        synchronized (key) {
            System.out.println( "current_key = " + key );
            if (this.redisTemplate.hasKey( key + "::accessToken" )) {
                accessToken = (String) this.redisTemplate.opsForValue().get( key + "::accessToken" );
                project_id = (String) this.redisTemplate.opsForValue().get( key + "::project_id" );
            } else {
                HashMap<String, String> map = GCloudUtils.generateGCloudJWTByJson( serviceAccountDir + "/"+key.replace( "|","/" ) + ".json" );
                accessToken = (String) map.get( "accessToken" );
                project_id = (String) map.get( "project_id" );
                this.redisTemplate.opsForValue().set( key + "::accessToken", accessToken, serviceAccountExpiration, TimeUnit.MINUTES );
                this.redisTemplate.opsForValue().set( key + "::project_id", project_id, serviceAccountExpiration, TimeUnit.MINUTES );
            }

        }

        //vertex-ai Claude不能在请求体里面带模型名称
        jsonObject.set( "model",null );
        String url = "https://" + region + "-aiplatform.googleapis.com/v1/projects/" + project_id + "/locations/" + region + "/publishers/anthropic/models/" + model + ":streamRawPredict";
        OkHttpClient.Builder clientBuilder = new OkHttpClient.Builder();
        okhttp3.RequestBody body = okhttp3.RequestBody.create(okhttp3.MediaType.get("application/json; charset=utf-8"), jsonObject.toString());
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + accessToken)
                .build();
        System.out.println("---------------------流式传输开始---------------------");

        HashMap<String,Integer> status =new HashMap<>(  );
        SseEmitter sseEmitter = new SseEmitter(0L);
        asyncService.asyncClaude(clientBuilder, request, sseEmitter,status);
        //阻塞状态来判断具体是开流还是返回错误
        while (true){
            //System.out.println( "status = " + status );
            if (status.containsKey( "status" )){

                if(status.get( "status" )==1){
                    response.addHeader("X-Content-Type-Options", "nosniff");
                    response.addHeader("X-Frame-Options", "SAMEORIGIN");
                    response.addHeader("X-XSS-Protection", "no");
                    response.addHeader("Content-Disposition", "attachment");
                    return sseEmitter;
                }else {
                    response.setStatus( 429 );
                    JSONObject entries = new JSONObject();
                    jsonObject.put( "code", status.getOrDefault( "code", 429 ));
                    return entries;

                }
            }
            try {
                Thread.sleep( 20 );
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

        }


    }

    /**
     * 处理非流式对话
     * @param response
     * @param key
     * @param jsonObject
     * @param model
     * @return
     * @throws IOException
     */
    public String claudeMessageWithNotStream(HttpServletResponse response,String key, JSONObject jsonObject,String model) throws IOException {
        jsonObject.put( "anthropic_version", "vertex-2023-10-16" );

        //目前只支持sonnet
        //用来拼接请求头
        String region = null;
        if (model.contains( "claude-3-5-sonnet" )){
            switch (RandomUtil.randomInt( 1,3 )){
                case 1:
                    region="europe-west1";
                    break;
                case 2:
                    region="us-east5";
                    break;
                default:
                    region="europe-west1";

            }
        }else if (model.contains("claude-3-opus")){
            switch (RandomUtil.randomInt( 1,3 )){
                case 1:
                    region="europe-west1";
                    break;
                case 2:
                    region="us-east5";
                    break;
                default:
                    region="europe-west1";
            }
        }

        //下面代码用来调用服务获取accessToken和project_id
        String accessToken = null;
        String project_id = null;
        synchronized (key) {
            System.out.println( "current_key = " + key );
            if (this.redisTemplate.hasKey( key + "::accessToken" )) {
                accessToken = (String) this.redisTemplate.opsForValue().get( key + "::accessToken" );
                project_id = (String) this.redisTemplate.opsForValue().get( key + "::project_id" );
            } else {
                HashMap<String, String> map = GCloudUtils.generateGCloudJWTByJson( serviceAccountDir +"/"+ key.replace( "|","/" ) + ".json" );
                accessToken = (String) map.get( "accessToken" );
                project_id = (String) map.get( "project_id" );
                this.redisTemplate.opsForValue().set( key + "::accessToken", accessToken, serviceAccountExpiration, TimeUnit.MINUTES );
                this.redisTemplate.opsForValue().set( key + "::project_id", project_id, serviceAccountExpiration, TimeUnit.MINUTES );
            }

        }
        //vertex-ai Claude不能在请求体里面带模型名称
        jsonObject.set( "model",null );
        String url = "https://" + region + "-aiplatform.googleapis.com/v1/projects/" + project_id + "/locations/" + region + "/publishers/anthropic/models/" + model + ":rawPredict";
        HttpRequest request = HttpUtil.createPost( url ).header( "Content-Type", "application/json" ).header( "charset", "utf-8" ).header( "Authorization", "Bearer " + accessToken ).body( jsonObject.toString()) ;
        request.setConnectionTimeout( 999999 );
        request.setReadTimeout( 999999 );
        HttpResponse execute = null;
        try {
            execute = request.execute();
        } catch (Exception var10) {
            var10.printStackTrace();
            throw new RJException( "请求谷歌服务器失败", -1 );
        }
        String body = execute.body();
        body=body.replace("msg_vrtx_","msg_");
        int status = execute.getStatus();
        response.setStatus( status );;

        System.out.println( region+"--->execute = " + body );
        if (status==200 && body.contains("id")){
            return  body ;
        }else {
            response.setStatus( 429 );;
            HashMap<Object, Object> map = new HashMap<>();
            map.put( "code", status);
            return JSONUtil.toJsonStr(map);
        }
    }
}
